package com.hzg.stream.function;

import cn.hutool.core.io.IoUtil;
import com.alibaba.druid.pool.DruidDataSource;
import com.hzg.util.CommonUtil;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.AsyncDataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.async.AsyncFunction;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import sun.nio.ch.IOUtil;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.Collections;
import java.util.concurrent.*;
import java.util.function.Supplier;


/**
 * 异步调用方法
 *
 * @Author: huangzhigao
 * @Date: 2022/4/19 17:01
 */
public class AsyncFunctionDemo {
    public static void main(String[] args) {
        StreamExecutionEnvironment env = CommonUtil.getEnv();
        DataStreamSource<Long> source = env.generateSequence(1, 10);
        //orderedWait数据输出的顺序与数据的进来顺序一致，不会出现乱序，unorderedWait会出现乱序
        AsyncDataStream.orderedWait(source,
                new MysqlAsyncFunction(),
                1000,  //超时时间
                TimeUnit.MILLISECONDS,
                1000); //可接受的异步请求次数
    }

    public static class MysqlAsyncFunction extends RichAsyncFunction<Long, Tuple2<Long, String>> {
        private transient ExecutorService executorService;
        private transient DruidDataSource dataSource = null;

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            executorService = Executors.newFixedThreadPool(10);
            //创建连接池（异步I/O 一个请求就是一个线程，一个请求对应一个连接）
            dataSource = new DruidDataSource();
            dataSource.setDriverClassName("com.mysql.jdbc.Driver");
            dataSource.setUsername("root");
            dataSource.setPassword("123456");
            dataSource.setUrl("jdbc:mysql://localhost:3306/day01?characterEncoding=utf8");
            dataSource.setMaxActive(10);
        }

        @Override
        public void asyncInvoke(Long input, ResultFuture<Tuple2<Long, String>> resultFuture) throws Exception {
            Future<String> future = executorService.submit(new Callable<String>() { // Callable是有返回值的submit
                @Override
                public String call() throws Exception {
                    return queryFromId(input);
                }
            });

            CompletableFuture<String> cf = CompletableFuture.supplyAsync(new Supplier<String>() {
                @Override
                public String get() {
                    try {
                        return future.get();
                    } catch (Exception e) {
                        return null;
                    }
                }
            });
            // 返回最终的resultFuture给方法
            cf.thenAccept((String result) -> {
                resultFuture.complete(Collections.singleton(Tuple2.of(input, result)));
            });

        }


        public String queryFromId(long id) throws SQLException {
            String sql = "SELECT id, name FROM orde WHERE id = ?";
            String result = null;
            Connection connection = null;
            PreparedStatement stmt = null;
            ResultSet rs = null;
            try {
                connection = dataSource.getConnection();
                stmt = connection.prepareStatement(sql);
                stmt.setLong(1, id);
                rs = stmt.executeQuery();
                while (rs.next()) {
                    result = rs.getString("name");
                    break;
                }
            } finally {
                IoUtil.close(rs);
                IoUtil.close(stmt);
                IoUtil.close(connection);
            }
            return result;

        }
    }
}
